'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 10_study_rerank_search.py
* @Time: 2025/9/24
* @All Rights Reserve By Brtc
'''
import os

import dotenv
import weaviate
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CohereRerank
from langchain_openai import OpenAIEmbeddings
from langchain_weaviate import WeaviateVectorStore
from weaviate.auth import AuthApiKey

dotenv.load_dotenv()
#1、创建向量数据库
embedding = OpenAIEmbeddings(model="text-embedding-3-small")

client = weaviate.connect_to_weaviate_cloud(
    skip_init_checks=True,
    cluster_url=os.getenv("WAEVIATE_URL"),
    auth_credentials=AuthApiKey(os.getenv("WEAVIATE_KEY"))
)
db = WeaviateVectorStore(
    client=client,
    index_name="TestParent",
    text_key = "text",
    embedding=OpenAIEmbeddings(model="text-embedding-3-small")
)

rerank = CohereRerank(model="rerank-multilingual-v3.0")

# 2构建压缩检索器
retriever = ContextualCompressionRetriever(
    base_retriever=db.as_retriever(),
    base_compressor=rerank
)

#3、执行搜索并排序
search_docs = retriever.invoke("关于LLMOPS应用配置的信息有那些？")
print(search_docs)
for one in search_docs:
    print(one)
print(len(search_docs))
client.close()

